from autom.config import cuda_if,cuda_num
import torch as t
from autom import postpro
import numpy as np
import scipy
from autom.siren import SirenNet

act_list = ['relu','tanh','siren','matrix']

def loss(E1,E2,M1,M2,img):
    return t.mean((E1*M1+E2*M2-img)**2)


def numpy2tensor(arr,outputer=None):
    if type(arr) == list:
        tensor = []
        for arr_now in arr:
            if type(arr_now) == np.ndarray:
                tensor_now = t.tensor(arr_now)
            else:
                tensor_now = arr_now
            if cuda_if:
                tensor_now = tensor_now.cuda(cuda_num)
            tensor.append(tensor_now)
    else:
        if type(arr) == np.ndarray:
            tensor = t.tensor(arr)
        else:
            tensor = arr
        if cuda_if:
            tensor = tensor.cuda(cuda_num)
    return tensor

def search_mask(extended1,extended2,img,mode='matrix',thre1=0.3,thre2=0.3,update_cor=True,
                searched_r=None,outputer=None,total_step=5000,progress=None):
    # TODO 统一mask的参数更新模式，把是否更新坐标作为一个选项，mode为matrix,inr,polynomial等
    extended1,extended2,img = numpy2tensor([extended1,extended2,img])
    searched_r = numpy2tensor(searched_r,outputer)
    M1,M2,c1,c2,r = main_opt(extended1,extended2,img,img_mode='tree',thre1=thre1,thre2=thre2,
                            update_cor=update_cor,mode=mode,searched_r=searched_r,
                            outputer=outputer,total_step=total_step,progress=progress)
    return M1,M2,c1,c2,r
    
def main_opt(extended1,extended2,img,img_mode='cor',thre1=0.3,thre2=0.3,update_cor=False,
            mode='matrix',searched_r=None,outputer=None,total_step=5000,progress=None):
    if type(searched_r) == type(None):
        r,cor1,cor2 = search_r(extended1,extended2,img,img_mode,thre1=thre1,thre2=thre2)
    else:
        r,cor1,cor2 = searched_r
    height,width = img.shape
    x_in = get_cor(height,width)
    # 同时优化坐标和mask
    imgt = t.tensor(img)
    my_mask = mask_class(mode=mode,img_shape=img.shape)
    if cuda_if:
        imgt = imgt.cuda(cuda_num)
        x_in = x_in.cuda(cuda_num)
    if update_cor:
        cor1 = t.nn.parameter.Parameter(cor1)
        cor2 = t.nn.parameter.Parameter(cor2)
        opt_c = t.optim.Adam([cor1,cor2],lr=1e-6)
    index1 = get_index(cor1,x_in)
    index2 = get_index(cor2,x_in)
    for i in range(total_step):
        if update_cor:
            opt_c.zero_grad()
        my_mask.update_M(update_mode='zero_grad')
        if img_mode == 'cor':
            E1 = cor2img(cor1,r,x_in,height,width)
            E2 = cor2img(cor2,r,x_in,height,width)
        else:
            E1 = tree2img(cor1,r,x_in,height,width,index1)
            E2 = tree2img(cor2,r,x_in,height,width,index2)
        loss_now = loss(E1,E2,my_mask.M1,my_mask.M2,imgt)
        loss_now.backward()
        my_mask.update_M(update_mode='step')
        if type(progress) != type(None):
            progress.value = int(i/total_step*100)
        if update_cor:
            opt_c.step()
    return my_mask.M1.detach().cpu().numpy(),my_mask.M2.detach().cpu().numpy(),cor1.detach().cpu().numpy(),cor2.detach().cpu().numpy(),r.detach().cpu().numpy()

class mask_class(object):
    def __init__(self,mode='matrix',img_shape=None):
        self.img_shape = img_shape
        self.mode = mode
        print(mode)
        if mode == 'matrix':
            self.matrix()
        elif mode == 'curv':
            self.curv()
        elif mode in act_list:
            self.inr(act=mode)
        else:
            raise('Wrong mode='+mode)
    
    def matrix(self):
        M1 = t.ones(self.img_shape)+t.rand(self.img_shape)
        M2 = t.ones(self.img_shape)+t.rand(self.img_shape)
        if cuda_if:
            M1 = M1.cuda(cuda_num)
            M2 = M2.cuda(cuda_num)
        self.M1 = t.nn.parameter.Parameter(M1)
        self.M2 = t.nn.parameter.Parameter(M2)
        self.opt = t.optim.Adam([self.M1,self.M2],lr=1e-3)

    def inr(self,act='siren'):
        if act == 'siren':
            self.net1 = SirenNet(2,256,1,2,w0_initial = 1.)
            self.net2 = SirenNet(2,256,1,2,w0_initial = 1.)
        elif act == 'relu':
            self.net1 = self.relu()
            self.net2 = self.relu()
        elif act == 'tanh':
            self.net1 = self.tanh()
            self.net2 = self.tanh()
        self.scale = 16
        self.x_in = get_cor(self.img_shape[0]//self.scale,self.img_shape[1]//self.scale)
        self.upsampled = t.nn.Upsample(size=self.img_shape,mode='bicubic')
        if cuda_if:
            self.x_in = self.x_in.cuda(cuda_num)
            self.net1 = self.net1.cuda(cuda_num)
            self.net2 = self.net2.cuda(cuda_num)
        self.opt1 = t.optim.Adam(self.net1.parameters(),lr=1e-4)
        self.opt2 = t.optim.Adam(self.net2.parameters(),lr=1e-4)
        self.M1 = self.net1(self.x_in).reshape(1,1,self.img_shape[0]//self.scale,self.img_shape[1]//self.scale)
        self.M1 = self.upsampled(self.M1).reshape(self.img_shape[0],self.img_shape[1])
        self.M2 = self.net2(self.x_in).reshape(1,1,self.img_shape[0]//self.scale,self.img_shape[1]//self.scale)
        self.M2 = self.upsampled(self.M2).reshape(self.img_shape[0],self.img_shape[1])

    def relu(self):
        return t.nn.Sequential(
            t.nn.Linear(2,64),
            t.nn.ReLU(),
            t.nn.Linear(64,64),
            t.nn.ReLU(),
            t.nn.Linear(64,64),
            t.nn.ReLU(),
            t.nn.Linear(64,64),
            t.nn.ReLU(),
            t.nn.Linear(64,64),
            t.nn.ReLU(),
            t.nn.Linear(64,1)
        )

    def tanh(self):
        return t.nn.Sequential(
            t.nn.Linear(2,64),
            t.nn.Tanh(),
            t.nn.Linear(64,64),
            t.nn.Tanh(),
            t.nn.Linear(64,64),
            t.nn.Tanh(),
            t.nn.Linear(64,64),
            t.nn.Tanh(),
            t.nn.Linear(64,64),
            t.nn.Tanh(),
            t.nn.Linear(64,1)
        )

    def dip(self):
        pass

    def curv(self):
        pass

    def update_M(self,update_mode='zero_grad'):
        if self.mode == 'matrix':
            if update_mode == 'zero_grad':
                self.opt.zero_grad()
            elif update_mode == 'step':
                self.opt.step()
            else:
                raise('Wrong update_mode='+update_mode)
        elif self.mode in act_list:
            if update_mode == 'zero_grad':
                self.opt1.zero_grad()
                self.opt2.zero_grad()
                self.M1 = self.net1(self.x_in).reshape(1,1,self.img_shape[0]//self.scale,self.img_shape[1]//self.scale)
                self.M1 = self.upsampled(self.M1).reshape(self.img_shape[0],self.img_shape[1])
                self.M2 = self.net2(self.x_in).reshape(1,1,self.img_shape[0]//self.scale,self.img_shape[1]//self.scale)
                self.M2 = self.upsampled(self.M2).reshape(self.img_shape[0],self.img_shape[1])
            elif update_mode == 'step':
                self.opt1.step()
                self.opt2.step()
            else:
                raise('Wrong update_mode='+update_mode)
        else:
            raise('Wrong mode = '+self.mode)
        


def curve2m(c,height,width):
    c.shape[0]
    pass

def search_mask_curve(extended1,extended2,img,order):
    # here exist order+1 coefficients
    c1 = t.rand(order+1)*1e-3
    c2 = t.rand(order+1)*1e-3
    E1 = t.tensor(extended1)
    E2 = t.tensor(extended2)
    imgt = t.tensor(img)
    if cuda_if:
        c1 = c1.cuda(cuda_num)
        c2 = c2.cuda(cuda_num)
        E1 = E1.cuda(cuda_num)
        E2 = E2.cuda(cuda_num)
        imgt = imgt.cuda(cuda_num)
    c1 = t.nn.parameter.Parameter(c1)
    c2 = t.nn.parameter.Parameter(c2)
    opt = t.optim.Adam([c1,c2],lr=1e-3)
    for _ in range(5000):
        opt.zero_grad()
        loss_now = loss(E1,E2,M1,M2,imgt)
        loss_now.backward()
        opt.step()
    pass






def cor2img(cor,r,x,height,width):
    # 从坐标映射到图像
    # 输入坐标cor \in N\times 2, 高斯基半径r, 查询坐标x \in M\times 2, 
    # 得到高斯基函数插值得到的x对应的像素值 f(x) = \sum_i e^{-(x-cor_i)^2/r} \in M\times 1
    # 优化计算量 ，将高斯基函数变成最近邻k个值加权求和 f(x) = \sum_{i\in N_x} ...    N_x为x的邻居集合 
    # 利用 kdtree来做，传入的参量也对应的发生了改变，变成了由cor形成的tree，每次只需要查询，不需要生成tree
    # 输出f(x).reshape(height,width)
    cor = cor.unsqueeze(2)
    x = x.T.unsqueeze(0)
    fx = t.zeros(height*width)
    if cuda_if:
        fx = fx.cuda(cuda_num)
    for i in range(height):
        for j in range(int(np.sqrt(width))):
            index_start = i*width+j*int(np.sqrt(width))
            index_end = i*width+(j+1)*int(np.sqrt(width))
            gx = t.exp(-t.sum((cor-x[:,:,index_start:index_end])**2,dim=1)/r)
            fx[index_start:index_end] = t.sum(gx,dim=0)
    return fx.reshape(height,width)

def cortree2img(height,width,coordinates,r):
    x = get_cor(height,width)
    coordinates,r,x = numpy2tensor([coordinates,r,x])
    index = get_index(coordinates,x)
    return tree2img(coordinates,r,x,height,width,index)


def get_index(cor,x):
    # 使用kdtree得到x距cor中最近的index
    # 由于cor后期只是微调，x是固定的，那么认为得到的index保持不变
    tree = scipy.spatial.cKDTree(cor.detach().cpu().numpy())
    _,index = tree.query(x.detach().cpu().numpy()) # dis,index \in \mathbb{R}^{x.shape[0]}
    return index


def tree2img(cor,r,x,height,width,index):
    # 从坐标形成的tree和输入x映射到图像
    # 默认选取x最近的一个点x'，若两者距离为d_{xx'}，那么计算得到x的像素值f(x) = e^{-d_{xx'}r}
    # 空间占用变少了，但是时间复杂度增高了
    # 由index计算xi隔得最近的距离，形成计算图
    x_n = cor[index,:]
    dis = t.sqrt(t.sum((x-x_n)**2,dim=1))
    fx = t.exp(-dis*r)
    return fx.reshape(height,width).T

def get_cor(height,width):
    x = np.linspace(0,width-1,width)
    y = np.linspace(0,height-1,height)
    xx,yy = np.meshgrid(x,y)
    xyz = np.stack([xx,yy],axis=2).astype('float32')
    in_put = t.tensor(xyz).reshape(-1,2)
    return in_put

def extract_atom(img,thre,sigma,r=1):
    coordinates = t.tensor(postpro.get_coor(img,thre=thre))
    x_in = get_cor(img.shape[0],img.shape[1])
    if cuda_if:
        coordinates = coordinates.cuda(cuda_num)
        x_in = x_in.cuda(cuda_num)
    index = get_index(coordinates,x_in)
    pre = tree2img(coordinates,r,x_in,img.shape[0],img.shape[1],index)
    return pre


def search_r(extended1,extended2,img,img_mode='cor',thre1=0.3,thre2=0.3,return_img=False,searched_r = None,progress=None):
    # 半径搜索
    # step1: 由拓展图像extendedi得到坐标
    # step2: 由坐标生成图像corimgi，先优化高斯基参数r使得corimgi和extendedi相匹配
    # step3: 基于高斯基参数大小不变，关键是如何构造损失函数使得坐标是可以微分的？ 高斯基函数累加!
    height,width = img.shape
    coordinates1 = t.tensor(postpro.get_coor(extended1,thre=thre1))
    coordinates2 = t.tensor(postpro.get_coor(extended2,thre=thre2))
    extended1 = t.tensor(extended1)
    extended2 = t.tensor(extended2)
    r = t.tensor([0.001])
    if searched_r != None:
        r = t.tensor(searched_r)

    x_in = get_cor(height,width)
    if cuda_if:
        r = r.cuda(cuda_num)
        x_in = x_in.cuda(cuda_num)
        extended1 = extended1.cuda(cuda_num)
        extended2 = extended2.cuda(cuda_num)
        coordinates1 = coordinates1.cuda(cuda_num)
        coordinates2 = coordinates2.cuda(cuda_num)
    r = t.nn.parameter.Parameter(r)
    opt_r = t.optim.Adam([r],lr=1e-3)
    index1 = get_index(coordinates1,x_in)
    index2 = get_index(coordinates2,x_in)
    if img_mode == 'cor':
        pre1 = cor2img(coordinates1,r,x_in,height,width)
        pre2 = cor2img(coordinates2,r,x_in,height,width)
    else:
        pre1 = tree2img(coordinates1,r,x_in,height,width,index1)
        pre2 = tree2img(coordinates2,r,x_in,height,width,index2)
    loss_r = t.mean((pre1-extended1)**2+(pre2-extended2)**2)
    if searched_r == None:
        for step_now in range(1000):
            opt_r.zero_grad()
            if img_mode == 'cor':
                pre1 = cor2img(coordinates1,r,x_in,height,width)
                pre2 = cor2img(coordinates2,r,x_in,height,width)
            else:
                pre1 = tree2img(coordinates1,r,x_in,height,width,index1)
                pre2 = tree2img(coordinates2,r,x_in,height,width,index2)
            loss_r = t.mean((pre1-extended1)**2+(pre2-extended2)**2)
            loss_r.backward()
            opt_r.step()
            if type(progress) != type(None):
                progress.value = int(step_now/10)
    if return_img:
        if searched_r == None:
            return r.detach().cpu().numpy(),pre1.detach().cpu().numpy(),pre2.detach().cpu().numpy(),coordinates1,coordinates2
        else:
            return pre1.detach().cpu().numpy(),pre2.detach().cpu().numpy(),coordinates1,coordinates2
    else:
        return r,coordinates1,coordinates2

def extend_correct(extend):

    pass



